// prefetcher.bpf.c
#include "vmlinux.h"
#include <bpf/bpf_helpers.h>
#include <bpf/bpf_tracing.h>

#include "prefetcher.h"

char LICENSE[] SEC("license") = "GPL";

// --- Maps (保持不变) ---
struct {
    __uint(type, BPF_MAP_TYPE_RINGBUF);
    __uint(max_entries, 256 * 1024);
} events SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_HASH);
    __uint(max_entries, 10240);
    __type(key, __u64);
    __type(value, struct access_history);
} history SEC(".maps");

struct {
    __uint(type, BPF_MAP_TYPE_ARRAY);
    __uint(max_entries, 1);
    __type(key, __u32);
    __type(value, struct config);
} config_map SEC(".maps");


// --- Context Struct (保持不变) ---
struct page_fault_ctx {
    unsigned short common_type;
    unsigned char common_flags;
    unsigned char common_preempt_count;
    int common_pid;
    unsigned long address;
    unsigned long ip;
    unsigned long error_code;
};

#define PAGE_SIZE 4096
#define CONFIDENCE_THRESHOLD 3
#define PREFETCH_PAGES 16

SEC("tracepoint/exceptions/page_fault_user")
int handle_page_fault(struct page_fault_ctx *ctx) {
    // --- 核心修正：回归最简单、最受验证器信任的模式 ---

    // 1. 获取当前进程/线程信息
    __u64 pid_tgid = bpf_get_current_pid_tgid();
    __u32 pid = pid_tgid >> 32;
    __u32 tid = (__u32)pid_tgid;

    // 2. 读取配置
    const __u32 key = 0;
    struct config *cfg;
    cfg = bpf_map_lookup_elem(&config_map, &key);
    if (!cfg) {
        // 如果配置 map 查找失败，这不应该发生，但作为保护直接退出
        return 0;
    }

    // 3. 执行过滤 (这是最关键的修改)
    // 直接解引用 cfg->target_pid。在检查 cfg != NULL 后，这是安全的。
    // 如果 cfg->target_pid 为 0，则条件 `pid != 0` 永不为真，继续执行。
    // 如果 cfg->target_pid 不为 0，则比较 pid。
    if (pid != cfg->target_pid && cfg->target_pid != 0) {
        return 0;
    }

    // 4. 获取缺页地址
    unsigned long addr = ctx->address;
    addr &= ~(PAGE_SIZE - 1);

    // 5. 查找或创建历史记录
    struct access_history *hist;
    hist = bpf_map_lookup_elem(&history, &pid_tgid);
    if (!hist) {
        // 使用零值初始化器，这是最安全的方式
        struct access_history new_hist = {};
        new_hist.last_addr = addr;
        bpf_map_update_elem(&history, &pid_tgid, &new_hist, BPF_ANY);
        return 0;
    }

    // 6. 模式识别
    if (hist->last_addr == 0) {
        hist->last_addr = addr;
        return 0;
    }
    
    long current_stride = addr - hist->last_addr;
    
    if (current_stride != 0 && current_stride == hist->last_stride) {
        hist->confidence++;
    } else {
        hist->confidence = 1;
    }

    hist->last_stride = current_stride;
    hist->last_addr = addr;

    // 7. 发送事件
    if (hist->confidence >= CONFIDENCE_THRESHOLD) {
        hist->confidence = 0; 
        
        struct pattern_info *e;
        e = bpf_ringbuf_reserve(&events, sizeof(*e), 0);
        if (!e) {
            return 0;
        }

        e->pid = pid;
        e->tid = tid;
        e->start_addr = addr + current_stride;
        e->stride = current_stride;
        e->size = PREFETCH_PAGES * PAGE_SIZE;
        
        bpf_ringbuf_submit(e, 0);
    }

    return 0;
}